from channels.generic.websocket import AsyncWebsocketConsumer
import json
import datetime as dt

class MarketDataConsumer(AsyncWebsocketConsumer):
    async def connect(self):
        # 接受WebSocket连接
        await self.accept()
        
        # 加入市场数据更新组
        await self.channel_layer.group_add(
            "market_data_updates",
            self.channel_name
        )
        
        print(f"WebSocket连接已建立，加入market_data_updates组 time = {dt.datetime.now()}")
        # print(f"当前channel_layer配置: {self.channel_layer}")
        # print(f"当前channel_name: {self.channel_name}")

    async def disconnect(self, close_code):
        # 离开市场数据更新组
        await self.channel_layer.group_discard(
            "market_data_updates",
            self.channel_name
        )
        print(f"WebSocket连接已断开，离开market_data_updates组")

    async def receive(self, text_data):
        try:
            data = json.loads(text_data)
            
            # 处理心跳消息
            if data.get('type') == 'heartbeat':
                await self.send(text_data=json.dumps({
                    'type': 'heartbeat',
                    'status': 'alive'
                }))
                return
                
            # 处理加入组请求
            if data.get('action') == 'join_group':
                group_name = data.get('group')
                if group_name:
                    await self.channel_layer.group_add(
                        group_name,
                        self.channel_name
                    )
                    await self.send(text_data=json.dumps({
                        'type': 'system',
                        'message': f'已加入组: {group_name}'
                    }))
                    
        except json.JSONDecodeError as e:
            await self.send(text_data=json.dumps({
                'type': 'error',
                'message': f'无效的消息格式: {str(e)}'
            }))

    async def send_progress(self, event):
        """处理直接发送的消息"""
        await self.send(text_data=json.dumps({
            'type': 'task_completed',
            'message': event.get('message'),
            'task_id': event.get('task_id'),
            'status': event.get('status'),
            'should_refresh': event.get('should_refresh', True)
        }))

    async def task_completed(self, event):
        """处理任务完成通知"""
        # 同时支持直接发送和组播
        await self.send(text_data=json.dumps({
            'type': 'task_completed',
            'message': event['message'],
            'task_id': event['task_id'],
            'status': event['status'],
            'should_refresh': event['should_refresh']
        }))